{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from netCDF4 import Dataset\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "from matplotlib.cm import get_cmap\n",
    "#from mpl_toolkits.basemap import Basemap\n",
    "import matplotlib.gridspec as gridspec\n",
    "from wrf import getvar, interplevel, to_np, get_basemap, latlon_coords, destagger, CoordPair, vertcross, interpline, eth\n",
    "import glob\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "import cmocean\n",
    "import pyresample\n",
    "from scipy.ndimage import gaussian_filter\n",
    "import matplotlib.patches as mpatches\n",
    "from matplotlib.lines import Line2D\n",
    "import Nio\n",
    "import geopy.distance\n",
    "import copy\n",
    "from scipy import stats\n",
    "\n",
    "### Import the scripts needed for importing satellite data ###\n",
    "from functions_import import import_VNP02, import_cldprop_l2, import_cldprop_uc_l2, import_MOD021KM, import_MOD06_L2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = \"mynn\"\n",
    "dom1 = \"d01\"\n",
    "dom2 = \"d02\"\n",
    "active_dom = \"d02\"\n",
    "sens = \"no_fractional_seaice\"\n",
    "day = \"12\"\n",
    "date = \"0313\"\n",
    "#root = \"/glade/scratch/tjuliano/doe_comble/new_lam_runs/WRF/WRF-ASR-CAO/test/\"\n",
    "root = \"/glade/campaign/uwyo/wyom0122/lam/new_runs/\"\n",
    "cases = [\"031220_no_edmf\",\"031220_mynn_l3\",\"031220\",\"031220_ysu\"]\n",
    "rst = \"/RESTART_D\"\n",
    "if active_dom == \"d01\":\n",
    "    plt_titles = [r\"$\\Delta$x=3km, level 2.5\", \"$\\Delta$x=3km, level 3\", \"$\\Delta$x=3km, level 2.5 EDMF\", \"$\\Delta$x=3km, YSU\"]\n",
    "    plt_titles2 = [r\"$\\Delta$x=3km\"\"\\n\"\"Level 2.5\", \"$\\Delta$x=3km\"\"\\n\"\"Level 3\", \"$\\Delta$x=3km\"\"\\n\"\"Level 2.5 EDMF\", \"$\\Delta$x=3km\"\"\\n\"\"YSU\"]\n",
    "    dx = 3000.\n",
    "else:\n",
    "    plt_titles = [r\"$\\Delta$x=1km, level 2.5\", \"$\\Delta$x=1km, level 3\", \"$\\Delta$x=1km, level 2.5 EDMF\", \"$\\Delta$x=1km, YSU\"]\n",
    "    plt_titles2 = [r\"$\\Delta$x=1km\"\"\\n\"\"Level 2.5\", \"$\\Delta$x=1km\"\"\\n\"\"Level 3\", \"$\\Delta$x=1km\"\"\\n\"\"Level 2.5 EDMF\", \"$\\Delta$x=1km\"\"\\n\"\"YSU\"]\n",
    "    dx = 1000.\n",
    "savedir = \"figs_new_nz136/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "lwprng = np.arange(-30,80,5)\n",
    "times = [9,11,13,19,25,31,37,43,49]\n",
    "time_str = ['28_1030z','28_1130z','28_1230z','28_1530z','28_1830z','28_2130z','29_0030z','29_0330z','29_0630z']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "lev = 10\n",
    "skipx = 6\n",
    "skipy = 6\n",
    "skipx2 = skipx*3\n",
    "skipy2 = skipy*3\n",
    "scale = 50\n",
    "barbw = 0.04\n",
    "g = 9.81"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = Nio.open_file(root + cases[0] + \"/wrfinput_\" + dom1,format='netcdf')\n",
    "cosalpha = s1.variables['COSALPHA'][-1,:,:]\n",
    "sinalpha = s1.variables['SINALPHA'][-1,:,:]\n",
    "s1.close()\n",
    "\n",
    "s1 = Nio.open_file(root + cases[3] + \"/wrfinput_\" + dom2,format='netcdf')\n",
    "cosalpha2 = s1.variables['COSALPHA'][-1,:,:]\n",
    "sinalpha2 = s1.variables['SINALPHA'][-1,:,:]\n",
    "s1.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Fig 9: EDMF-ONLY\n",
    "## For NSF proposal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "files_wind0 = sorted(glob.glob(root + cases[0] + rst + \"/wrfout_wind_\" + active_dom + \"*\"))\n",
    "files_wind1 = sorted(glob.glob(root + cases[1] + rst + \"/wrfout_wind_\" + active_dom + \"*\"))\n",
    "files_wind2 = sorted(glob.glob(root + cases[2] + rst + \"/wrfout_wind_\" + active_dom + \"*\"))\n",
    "files_wind3 = sorted(glob.glob(root + cases[3] + rst + \"/wrfout_wind_\" + active_dom + \"*\"))\n",
    "\n",
    "files_cloud0 = sorted(glob.glob(root + cases[0] + rst + \"/wrfout_cloud_\" + active_dom + \"*\"))\n",
    "files_cloud1 = sorted(glob.glob(root + cases[1] + rst + \"/wrfout_cloud_\" + active_dom + \"*\"))\n",
    "files_cloud2 = sorted(glob.glob(root + cases[2] + rst + \"/wrfout_cloud_\" + active_dom + \"*\"))\n",
    "files_cloud3 = sorted(glob.glob(root + cases[3] + rst + \"/wrfout_cloud_\" + active_dom + \"*\"))\n",
    "\n",
    "files_rad0 = sorted(glob.glob(root + cases[0] + rst + \"/wrfout_turb_\" + active_dom + \"*\"))\n",
    "files_rad1 = sorted(glob.glob(root + cases[1] + rst + \"/wrfout_turb_\" + active_dom + \"*\"))\n",
    "files_rad2 = sorted(glob.glob(root + cases[2] + rst + \"/wrfout_turb_\" + active_dom + \"*\"))\n",
    "files_rad3 = sorted(glob.glob(root + cases[3] + rst + \"/wrfout_turb_\" + active_dom + \"*\"))\n",
    "\n",
    "times = [12]\n",
    "first_pass1 = 0\n",
    "zmax = 5550\n",
    "zint = 50\n",
    "\n",
    "qscale = 4.\n",
    "skipx = 5\n",
    "skipz = 3\n",
    "for i in times: #np.arange(12,len(files_cloud0),6):\n",
    "\tfor j in np.arange(2,3):\n",
    "\t\tprint ('Reading file...')\n",
    "\t\tif j == 0:\n",
    "\t\t\ts1 = Nio.open_file(files_wind0[i],format='netcdf')\n",
    "\t\t\ts2 = Nio.open_file(files_cloud0[i],format='netcdf')\n",
    "\t\t\ts3 = Nio.open_file(files_rad0[i],format='netcdf')\n",
    "\t\t\tprint (files_wind0[i])\n",
    "\t\telif j == 1:\n",
    "\t\t\ts1 = Nio.open_file(files_wind1[i],format='netcdf')\n",
    "\t\t\ts2 = Nio.open_file(files_cloud1[i],format='netcdf')\n",
    "\t\t\ts3 = Nio.open_file(files_rad1[i],format='netcdf')\n",
    "\t\t\tprint (files_wind1[i])\n",
    "\t\telif j == 2:\n",
    "\t\t\ts1 = Nio.open_file(files_wind2[i],format='netcdf')\n",
    "\t\t\ts2 = Nio.open_file(files_cloud2[i],format='netcdf')\n",
    "\t\t\ts3 = Nio.open_file(files_rad2[i],format='netcdf')\n",
    "\t\t\tprint (files_wind2[i])\n",
    "\t\telif j == 3:\n",
    "\t\t\ts1 = Nio.open_file(files_wind3[i],format='netcdf')\n",
    "\t\t\ts2 = Nio.open_file(files_cloud3[i],format='netcdf')\n",
    "\t\t\ts3 = Nio.open_file(files_rad3[i],format='netcdf')\n",
    "\t\t\tprint (files_wind3[i])\n",
    "\t\tif first_pass1 == 0:\n",
    "\t\t\ts4 = Nio.open_file(root + cases[0] + \"/wrfinput_\" + active_dom,format='netcdf')\n",
    "\t\t\tlat = getvar(s4,\"lat\",meta=False)\n",
    "\t\t\tlon = getvar(s4,\"lon\",meta=False)\n",
    "        \n",
    "\t\tif first_pass1 == 0:\n",
    "\t\t\tlon2d = s1.variables['XLONG'][-1,:,:]\n",
    "\t\t\tlat2d = s1.variables['XLAT'][-1,:,:]\n",
    "\t\t\tgrid = pyresample.geometry.GridDefinition(lats=lat2d, lons=lon2d)\n",
    "\n",
    "\t\t\tfirst_pass1 = 1\n",
    "\n",
    "\t\ttheta = getvar(s1,\"T\", meta=True)+300.\n",
    "\t\tw = getvar(s1,\"W\", meta=True)\n",
    "\t\tu = s1.variables['U'][-1,:,:,:]\n",
    "\t\tz = getvar(s1,\"z\", units=\"m\",meta=False)\n",
    "\t\tpblh = s1.variables['PBLH'][-1,:,:]\n",
    "\t\trho = 1./s1.variables['ALT'][-1,:,:,:]\n",
    "        \n",
    "\t\t# Create the start point and end point for the cross section\n",
    "\t\tstart_point = CoordPair(lat=69.4, lon=12.5)\n",
    "\t\tend_point = CoordPair(lat=69.4, lon=16.5)\n",
    "            \n",
    "\t\tvar_list = ['QKE','QSHEAR','QBUOY','QDISS','QWT','EDMF_A','EDMF_W']\n",
    "\t\tvar_list2 = ['A$_{EDMF}$','Q$_{C}$','CDNC','CCN']\n",
    "\t\tvar_list3 = ['[-]','[g kg$^{-1}$]','[cm$^{-3}$]','[cm$^{-3}$]']\n",
    "\t\tEDMF_A = getvar(s3,'EDMF_A', meta=False)\n",
    "\t\tEDMF_W = getvar(s3,'EDMF_W', meta=False)\n",
    "\t\tw = destagger(w,0)\n",
    "\t\tu = destagger(u,2)\n",
    "        \n",
    "\t\tqcloud = s2.variables['QCLOUD'][-1,:,:,:]*1000.\n",
    "\t\tqice = s2.variables['QICE'][-1,:,:,:]*1000.\n",
    "\t\tqsnow = s2.variables['QSNOW'][-1,:,:,:]*1000.\n",
    "\t\tqrain = s2.variables['QRAIN'][-1,:,:,:]*1000.\n",
    "\t\tqgraup = s2.variables['QGRAUP'][-1,:,:,:]*1000.\n",
    "\t\tqncloud = s2.variables['QNCLOUD'][-1,:,:,:]*rho*1.0e-6\n",
    "\t\tqnwfa = s2.variables['QNWFA'][-1,:,:,:]*rho*1.0e-6\n",
    "        \n",
    "\t\tqliq = qrain #qcloud+qrain\n",
    "\t\tqice = qice\n",
    "\t\tqfrz = qsnow+qgraup\n",
    "\n",
    "\t################################################################################################################\n",
    "\t############################################### LWP ############################################################\n",
    "\t################################################################################################################\n",
    "        \n",
    "\t\tfig = plt.figure(figsize=(15,20))\n",
    "\t\t#ax1 = fig.add_axes([0.7,0.05,0.05,0.4])\n",
    "        \n",
    "\t\t# Compute the vertical cross-section interpolation.  Also, include the\n",
    "\t\t# lat/lon points along the cross-section.\n",
    "\t\tpblh_cross = interpline(pblh, wrfin=s1, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "\t\ttheta_cross = vertcross(theta, z, levels=np.arange(0,zmax,zint), wrfin=s1, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "\t\tw_cross = vertcross(w, z, levels=np.arange(0,zmax,zint), wrfin=s1, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "\t\tu_cross = vertcross(u, z, levels=np.arange(0,zmax,zint), wrfin=s1, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "\t\tEDMF_A_cross = vertcross(EDMF_A, z, levels=np.arange(0,zmax,zint), wrfin=s1, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "\t\tEDMF_W_cross = vertcross(EDMF_W, z, levels=np.arange(0,zmax,zint), wrfin=s1, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "\t\tqliq_cross = vertcross(qliq, z, levels=np.arange(0,zmax,zint), wrfin=s1, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "\t\tqcloud_cross = vertcross(qcloud, z, levels=np.arange(0,zmax,zint), wrfin=s1, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "\t\tqncloud_cross = vertcross(qncloud, z, levels=np.arange(0,zmax,zint), wrfin=s1, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "\t\tqnwfa_cross = vertcross(qnwfa, z, levels=np.arange(0,zmax,zint), wrfin=s1, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "\t\tqice_cross = vertcross(qice, z, levels=np.arange(0,zmax,zint), wrfin=s1, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "\t\tqfrz_cross = vertcross(qfrz, z, levels=np.arange(0,zmax,zint), wrfin=s1, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "\t\tfor jj in np.arange(4):\n",
    "\t\t\tplt.subplot(4,1,jj+1)\n",
    "\t\t\trng = np.arange(268,282,1)\n",
    "\t\t\tif jj == 0:\n",
    "\t\t\t\tvar2plt = EDMF_A_cross\n",
    "\t\t\t\trng2 = np.arange(0,0.1,0.01)\n",
    "\t\t\t\tcmap=cmocean.cm.thermal\n",
    "\t\t\t\textend='max'\n",
    "\t\t\telif jj == 1:\n",
    "\t\t\t\tvar2plt = qcloud_cross\n",
    "\t\t\t\trng2 = np.arange(0,0.11,0.01)\n",
    "\t\t\t\tcmap=cmocean.cm.thermal\n",
    "\t\t\t\textend='max'\n",
    "\t\t\telif jj == 2:\n",
    "\t\t\t\tvar2plt = qncloud_cross\n",
    "\t\t\t\trng2 = np.arange(0,18,2)\n",
    "\t\t\t\tcmap=cmocean.cm.thermal\n",
    "\t\t\t\textend='max'\n",
    "\t\t\telif jj == 3:\n",
    "\t\t\t\tvar2plt = qnwfa_cross\n",
    "\t\t\t\trng2 = np.arange(0,52,4)\n",
    "\t\t\t\tcmap=cmocean.cm.oxy\n",
    "\t\t\t\textend='max'\n",
    "\t\t\tvar_plt_contourf = plt.contourf(to_np(var2plt[:,:]),rng2,cmap=cmap,extend=extend,zorder=2)\n",
    "#\t\t\ttheta_contourf = ax.contourf(to_np(theta_cross[:,:]),rng,cmap=cmocean.cm.thermal,extend='both',zorder=2)\n",
    "\t\t\ttheta_contour = plt.contour(to_np(theta_cross[:,:]),rng,colors='magenta',linewidths=3.,zorder=4)\n",
    "\t\t\tplt.clabel(theta_contour,fmt='%d',zorder=4)\n",
    "\t\t\tpblh_contour = plt.plot(np.arange(len(pblh_cross)),to_np(pblh_cross[:])/zint,c='white',lw=3.,zorder=4)\n",
    "\t\t\tif jj == 1:\n",
    "\t\t\t\twp_contour = plt.contour(to_np(w_cross[:,:]),[0.5],colors='limegreen',linewidths=2.,zorder=4)\n",
    "\t\t\t\twm_contour = plt.contour(to_np(w_cross[:,:]),[-0.5],colors='lightskyblue',linewidths=2.,linestyles='--',zorder=4)\n",
    "\t\t\tif jj == 0:\n",
    "\t\t\t\tqliq_contour = plt.contourf(to_np(qliq_cross[:,:]),[0.001,1],colors='none',hatches='.',zorder=7)\n",
    "\t\t\t\tfor ii, collection in enumerate(qliq_contour.collections):\n",
    "\t\t\t\t\tcollection.set_edgecolor('red')\n",
    "\t\t\t\t# Doing this also colors in the box around each level\n",
    "\t\t\t\t# We can remove the colored line around the levels by setting the linewidth to 0\n",
    "\t\t\t\tfor collection in qliq_contour.collections:\n",
    "\t\t\t\t\tcollection.set_linewidth(0.)\n",
    "\t\t\t\tqice_contour = plt.contourf(to_np(qice_cross[:,:]),[0.001,1],colors='none',hatches='o',zorder=6)\n",
    "\t\t\t\tfor ii, collection in enumerate(qice_contour.collections):\n",
    "\t\t\t\t\tcollection.set_edgecolor('limegreen')\n",
    "\t\t\t\t# Doing this also colors in the box around each level\n",
    "\t\t\t\t# We can remove the colored line around the levels by setting the linewidth to 0\n",
    "\t\t\t\tfor collection in qice_contour.collections:\n",
    "\t\t\t\t\tcollection.set_linewidth(0.)\n",
    "\t\t\t\tqfrz_contour = plt.contourf(to_np(qfrz_cross[:,:]),[0.1,1],colors='none',hatches='*',zorder=5)\n",
    "\t\t\t\tfor ii, collection in enumerate(qfrz_contour.collections):\n",
    "\t\t\t\t\tcollection.set_edgecolor('lightskyblue')\n",
    "\t\t\t\t# Doing this also colors in the box around each level\n",
    "\t\t\t\t# We can remove the colored line around the levels by setting the linewidth to 0\n",
    "\t\t\t\tfor collection in qfrz_contour.collections:\n",
    "\t\t\t\t\tcollection.set_linewidth(0.)\n",
    "\t\t\tif jj > 1:\n",
    "\t\t\t\txlonv = np.arange(np.shape(to_np(w_cross[:,:]))[1])\n",
    "\t\t\t\tzv = np.arange(np.shape(to_np(w_cross[:,:]))[0])\n",
    "\t\t\t\tmean_u = np.mean(to_np(u_cross),axis=1)\n",
    "\t\t\t\tpert_u = np.empty([np.shape(u_cross)[0],np.shape(u_cross)[1]])\n",
    "\t\t\t\tfor iii in np.arange(np.shape(u_cross)[0]):\n",
    "\t\t\t\t\tpert_u[iii,:] = u_cross[iii,:] - mean_u[iii]\n",
    "\t\t\t\tsmooth_u = gaussian_filter(to_np(pert_u),sigma=1) #smooth2d(pert_uv,2)\n",
    "\t\t\t\tsmooth_w = gaussian_filter(to_np(w_cross),sigma=1) #smooth2d(wv,2)\n",
    "\t\t\t\tquiv = plt.quiver(xlonv[::skipx],zv[::skipz],to_np(smooth_u[::skipz,::skipx]),qscale*to_np(smooth_w[::skipz,::skipx]),scale=100,width=0.003,facecolor='limegreen',edgecolor='k',lw=0.15,zorder=3)\n",
    "\t\t\tcoord_pairs = to_np(theta_cross.coords[\"xy_loc\"])\n",
    "\t\t\tx_ticks = np.arange(coord_pairs.shape[0])\n",
    "\t\t\tx_labels = [pair.latlon_str() for pair in to_np(coord_pairs)]\n",
    "\t\t\tif jj == 3:\n",
    "\t\t\t\tplt.xticks(x_ticks[::30],['0','30','60','90','120','150'], fontsize=18)\n",
    "\t\t\t\tplt.xlabel('Distance [km]',fontsize=18)\n",
    "\t\t\telse:\n",
    "\t\t\t\tplt.xticks(x_ticks[::30],[])\n",
    "\t\t\tplt.yticks(np.arange(0,np.shape(to_np(var2plt[:,:]))[0],10),['0','','1000','','2000','','3000','','4000','','5000',''],fontsize=16)\n",
    "\t\t\tplt.ylabel('Height [m]',fontsize=18)\n",
    "\t\t\tcbar = plt.colorbar(var_plt_contourf)\n",
    "\t\t\tcbar.ax.tick_params(labelsize=18)\n",
    "\t\t\tcbar.ax.set_ylabel(var_list3[jj], labelpad=40, rotation=270, size=16)\n",
    "\t\t\tplt.text(5,107,var_list2[jj],fontsize=22,bbox=dict(facecolor='darkgray', edgecolor='green', alpha=1.0, boxstyle='round'),zorder=5)\n",
    "            \n",
    "\t\t\tif jj == 3:\n",
    "\t\t\t\tqliq_artists = mpatches.Patch(facecolor='None',edgecolor='red',hatch='.',label='Q$_{r}$')\n",
    "\t\t\t\tqice_artists= mpatches.Patch(facecolor='None',edgecolor='limegreen',hatch='o',label='Q$_{i}$')\n",
    "\t\t\t\tqfrz_artists = mpatches.Patch(facecolor='None',edgecolor='lightskyblue',hatch='*',label='Q$_{s}$+Q$_{g}$')\n",
    "\t\t\t\twp_artists = Line2D([0, 1], [0, 1], color='limegreen', lw=3., label='W+')\n",
    "\t\t\t\twm_artists = Line2D([0, 1], [0, 1], color='lightskyblue', lw=3., ls='--', label='W-')\n",
    "\t\t\t\tfig.legend(handles=[qliq_artists,qice_artists,qfrz_artists,wp_artists,wm_artists],loc='lower center',bbox_to_anchor=(0.55,0.03),bbox_transform=plt.gcf().transFigure,ncol=5,fontsize=18)\n",
    "\t\t\t\tqk = plt.quiverkey(quiv, 0.15, 0.05, 4, '4 m/s', labelpos='E',coordinates='figure',fontproperties={'weight': 'bold','size':18},zorder=6)\n",
    "\n",
    "#plt.show()\n",
    "plt.savefig(savedir + \"WRF-EDMF-ONLY-XSECT-NSF.png\",dpi=600,bbox_inches='tight')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
